package main

import (
	"fmt"
	"math"
	"os"
	"strconv"
	"strings"
	"time"
)

// NPU 结构体定义
type NPU struct {
	ID               int
	Model            string
	Health           string
	Power            float64
	Temp             int
	AiCore           int
	MemoryUse        int
	MaxMemory        int
	MemoryUsePercent float64
}

// NPUAvg 结构体定义
type NPUAvg struct {
	Power            float64
	MaxMemory        int
	MemoryUse        int
	AiCore           int
	Health           string
	MemoryUsePercent float64
}

// Table 结构体定义
type Table struct {
	Title      string
	FieldNames []string
	Rows       [][]string
}

// 模拟 color_row 函数
func colorRow(npu NPU) []string {
	return []string{
		strconv.Itoa(npu.ID),
		strconv.Itoa(npu.Temp),
		fmt.Sprintf("%.1f", npu.Power),
		fmt.Sprintf("%d", npu.MemoryUse),
		strconv.Itoa(npu.AiCore),
		npu.Health,
		fmt.Sprintf("%.3f", npu.MemoryUsePercent),
		strconv.Itoa(npu.AiCore),
	}
}

// 模拟 color_str 函数
func colorStr(s, color string) string {
	return s
}

// 模拟 Table 的方法
func (t *Table) ClearRows() {
	t.Rows = [][]string{}
}

func (t *Table) AddRow(row []string) {
	t.Rows = append(t.Rows, row)
}

func convertPythonLogic(re []string, npuNum int, isWatch string) bool {
	var table Table
	npuMapping := make(map[int]*NPU)
	npuAvg := NPUAvg{Health: "OK"}

	// 检查标题是否包含 "npu"
	if strings.Contains(strings.ToLower(table.Title), "npu") {
		table.ClearRows()
	} else {
		if len(re) > 1 {
			versionParts := strings.Split(strings.Split(re[1], "|")[0], " ")
			if len(versionParts) > 2 {
				version := versionParts[2]
				table.Title = "npu-smi " + version
				table.FieldNames = []string{"NPU", "Temp(C)", "Pow(W)", " HBM-Usage(MB) ", "NPU-Util", " Health ", "Memory Utilization", "NPU Utilization"}
			}
		}
	}

	for i := 4; i < npuNum*2+4; i += 2 {
		tem := strings.Split(re[i], "|")
		tem0 := strings.Fields(tem[0])
		temID, _ := strconv.Atoi(tem0[0])
		if _, exists := npuMapping[temID]; !exists {
			npuMapping[temID] = &NPU{ID: temID}
		}
		npuMapping[temID].Model = tem0[1]
		tem1 := strings.Fields(tem[1])
		npuMapping[temID].Health = tem1[0]
		tem2 := strings.Fields(tem[2])
		npuMapping[temID].Power, _ = strconv.ParseFloat(tem2[0], 64)
		npuMapping[temID].Temp, _ = strconv.Atoi(tem2[1])

		tem = strings.Split(re[i+1], "|")
		tem2 = strings.Fields(strings.ReplaceAll(tem[2], "/", " "))
		npuMapping[temID].AiCore, _ = strconv.Atoi(tem2[0])
		npuMapping[temID].MemoryUse, _ = strconv.Atoi(tem2[3])
		npuMapping[temID].MaxMemory, _ = strconv.Atoi(tem2[4])
		npuMapping[temID].MemoryUsePercent = math.Round(float64(npuMapping[temID].MemoryUse)/float64(npuMapping[temID].MaxMemory)*1000) / 1000

		npuAvg.Power += npuMapping[temID].Power
		npuAvg.MaxMemory += npuMapping[temID].MaxMemory
		npuAvg.MemoryUse += npuMapping[temID].MemoryUse
		npuAvg.AiCore += npuMapping[temID].AiCore

		table.AddRow(colorRow(*npuMapping[temID]))
	}

	if isWatch == "watch" {
		fmt.Print("\033[H\033[1B\033[J")
		formattedTime := time.Now().Format("2006-01-02 15:04:05")
		fmt.Println(formattedTime + strings.Repeat(" ", 122) + "Press " + colorStr("q", "red") + " to quit")
	}

	npuAvg.Power = math.Round(npuAvg.Power*10) / 10
	npuAvg.MemoryUsePercent = math.Round(float64(npuAvg.MemoryUse)/float64(npuAvg.MaxMemory)*1000) / 1000
	npuAvg.MemoryUse = int(math.Round(float64(npuAvg.MemoryUse)))
	npuAvg.AiCore = int(math.Round(float64(npuAvg.AiCore) / 8))

	table.AddRow(colorRow(NPU{
		Power:            npuAvg.Power,
		MaxMemory:        npuAvg.MaxMemory,
		MemoryUse:        npuAvg.MemoryUse,
		AiCore:           npuAvg.AiCore,
		Health:           npuAvg.Health,
		MemoryUsePercent: npuAvg.MemoryUsePercent,
	}))

	fmt.Println(table)
	return true
}

func main0() {
	// 读取demo2.txt的内容
	fileContent, err := os.ReadFile("./gpu_demo/demo2.txt")
	if err != nil {
		fmt.Println("读取文件时出错：", err)
		return
	}
	re := []string{string(fileContent)}
	npuNum := 1
	isWatch := "watch"
	convertPythonLogic(re, npuNum, isWatch)
}
